import os
import sys
from util import utils
# =============predo=============== #
func = utils()
sys.path.append(func.getParentDirectoryPath())  # It has get the ParentFolder.
workspacePath = os.path.dirname(func.getParentDirectoryPath())
# ================================= #
# from ..model.XDUnet import XDUnet
# or
from model.XDUnet import XDUnet
from loadData import data

class train:
    models  = ['2DUnet','3DUnet']
    hyperParameters = {}

    def __init__(self,modelStr,taskNum,**Kwargs):
        assert (taskNum in [3,4]),'不符合'
        assert (modelStr in self.models),'训练模型不在列表'
        self.taskNum  = taskNum
        self.modelStr = modelStr
        self.hyperParameters = Kwargs
        self.requireParameters = ['rate','n_classes','topchannels','input_width','input_height','learn_rate','batch_size','max_epochs']

    def modelGet(self):
        dict_   = self.hyperParameters
        if(self.modelStr == self.models[0]):
            return XDUnet.Unet2D(dict_['topChannel'],dict_['n_classes'],dict_['input_width'],dict_['input_height'])
        
        if(self.modelStr == self.models[1]):
            return XDUnet.Unet3D(dict_['topChannel'],dict_['n_classes'],dict_['input_width'],dict_['input_height'])

    def trainMainCall(self):
        dataset = data(os.path.join(workspacePath,'dataset'),self.taskNum)
        assert (['rate','n_classes'] in self.hyperParameters.keys()),"输入超参数不足"
        x_train,x_valid,y_train,y_valid = dataset.trainTest(self.hyperParameters['rate'],self.hyperParameters['n_classes'])
        mainModel = self.modelGet()
        return mainModel


if __name__ == "__main__":
    a = train('2DUnet',3,n_classes=2,rate=0.1)     
    print(a.hyperParameters)
    a.trainMainCall()

